from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from webdriver_manager.firefox import GeckoDriverManager
from utils.config_utils import config
from utils.logger_utils import logger
import os

class DriverUtils:
    def __init__(self):
        self.browser = config.get_value('test_env', 'browser').lower()
        self.base_url = config.get_value('test_env', 'base_url')
        self.implicit_wait = config.get_int('test_env', 'implicit_wait')
        self.driver = None
    
    def get_driver(self):
        try:
            if self.browser == 'chrome':
                options = webdriver.ChromeOptions()
                options.add_argument('--start-maximized')
                options.add_argument('--disable-gpu')
                options.add_argument('--no-sandbox')
                options.add_argument("detach",True)
                # options.add_argument('--headless')  # 无头模式，可根据需要启用
                
                # 使用 webdriver_manager 自动管理驱动
                service = Service(ChromeDriverManager().install())
                self.driver = webdriver.Chrome(service=service, options=options)
            elif self.browser == 'firefox':
                options = webdriver.FirefoxOptions()
                options.add_argument('--start-maximized')
                
                service = Service(GeckoDriverManager().install())
                self.driver = webdriver.Firefox(service=service, options=options)
            else:
                logger.error(f"不支持的浏览器: {self.browser}")
                raise ValueError(f"不支持的浏览器: {self.browser}")
            
            self.driver.implicitly_wait(self.implicit_wait)
            self.driver.get(self.base_url)
            logger.info(f"成功启动 {self.browser} 浏览器，访问 {self.base_url}")
            return self.driver
        except Exception as e:
            logger.error(f"启动浏览器失败: {e}")
            raise
    
    def quit_driver(self):
        if self.driver:
            try:
                self.driver.quit()
                logger.info("关闭浏览器成功")
            except Exception as e:
                logger.error(f"关闭浏览器失败: {e}")
    
    def take_screenshot(self, file_name):
        if self.driver:
            screenshots_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../screenshots')
            if not os.path.exists(screenshots_dir):
                os.makedirs(screenshots_dir)
            
            file_path = os.path.join(screenshots_dir, f'{file_name}_{datetime.now().strftime("%Y%m%d_%H%M%S")}.png')
            try:
                self.driver.save_screenshot(file_path)
                logger.info(f"截图保存成功: {file_path}")
                return file_path
            except Exception as e:
                logger.error(f"截图失败: {e}")
                return None

# 单例模式
driver_utils = DriverUtils()  